import os
import cv2
import json
import numpy as np
from paddleocr import PaddleOCR
import layoutparser as lp
from pdf2image import convert_from_path

class PDFOCRProcessor:
    def __init__(self, lang='ch'):
        # 初始化OCR引擎
        self.ocr = PaddleOCR(use_angle_cls=True, lang=lang)
        
        # 初始化布局分析模型
        self.layout_model = lp.Detectron2LayoutModel(
            'lp://PubLayNet/faster_rcnn_R_50_FPN_3x/config',
            extra_config=["MODEL.ROI_HEADS.SCORE_THRESH_TEST", 0.5],
            label_map={0: "Text", 1: "Title", 2: "List", 3: "Table", 4: "Figure"}
        )
    
    def process_pdf(self, pdf_path, output_dir=None):
        """处理PDF文件，提取文本和布局信息"""
        if output_dir is None:
            output_dir = os.path.dirname(pdf_path)
        os.makedirs(output_dir, exist_ok=True)
        
        # 将PDF转换为图像
        images = convert_from_path(pdf_path)
        
        results = []
        for i, image in enumerate(images):
            # 转换为OpenCV格式
            img_np = np.array(image)
            img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
            
            # 保存图像用于处理
            img_path = os.path.join(output_dir, f"page_{i+1}.jpg")
            cv2.imwrite(img_path, img_np)
            
            # 处理单页
            page_result = self.process_page(img_path, i+1)
            results.append(page_result)
            
        # 保存结果为JSON
        output_json = os.path.join(output_dir, f"{os.path.basename(pdf_path).split('.')[0]}_result.json")
        with open(output_json, 'w', encoding='utf-8') as f:
            json.dump(results, f, ensure_ascii=False, indent=2)
            
        return output_json
    
    def process_page(self, img_path, page_num):
        """处理单个页面，提取文本和布局"""
        # 读取图像
        image = cv2.imread(img_path)
        
        # OCR识别
        ocr_result = self.ocr.ocr(img_path, cls=True)
        
        # 布局分析
        layout_image = lp.load_image(img_path)
        layout = self.layout_model.detect(layout_image)
        
        # 构建结果
        page_result = {
            "page": page_num,
            "blocks": []
        }
        
        # 处理布局元素
        for block in layout:
            block_coords = block.coordinates
            x1, y1, x2, y2 = [int(coord) for coord in [block_coords[0], block_coords[1], block_coords[2], block_coords[3]]]
            
            # 提取该区域内的OCR结果
            block_texts = []
            for line in ocr_result[0]:
                bbox = line[0]
                text = line[1][0]
                confidence = line[1][1]
                
                # 检查文本框是否在当前布局块内
                text_x1, text_y1 = bbox[0][0], bbox[0][1]
                text_x2, text_y2 = bbox[2][0], bbox[2][1]
                
                if (x1 <= text_x1 <= x2 and y1 <= text_y1 <= y2) or \
                   (x1 <= text_x2 <= x2 and y1 <= text_y2 <= y2):
                    block_texts.append({
                        "text": text,
                        "confidence": float(confidence),
                        "bbox": [[float(p[0]), float(p[1])] for p in bbox]
                    })
            
            # 根据布局类型处理
            if block.type == "Table":
                # 对于表格，尝试提取行和列
                table_content = self.extract_table_structure(image[y1:y2, x1:x2], block_texts)
                page_result["blocks"].append({
                    "type": "table",
                    "coordinates": [x1, y1, x2, y2],
                    "content": table_content
                })
            else:
                # 对于其他类型的块
                page_result["blocks"].append({
                    "type": block.type.lower(),
                    "coordinates": [x1, y1, x2, y2],
                    "content": block_texts
                })
        
        return page_result
    
    def extract_table_structure(self, table_image, texts):
        """尝试提取表格结构"""
        # 这里可以实现更复杂的表格结构提取算法
        # 简单实现：根据文本Y坐标分组为行
        if not texts:
            return []
            
        # 按Y坐标排序
        texts.sort(key=lambda x: x["bbox"][0][1])
        
        # 简单的行分组（可以改进）
        rows = []
        current_row = [texts[0]]
        row_y = texts[0]["bbox"][0][1]
        
        for text in texts[1:]:
            text_y = text["bbox"][0][1]
            # 如果Y坐标相差不大，认为是同一行
            if abs(text_y - row_y) < 20:  # 阈值可调整
                current_row.append(text)
            else:
                # 按X坐标排序当前行
                current_row.sort(key=lambda x: x["bbox"][0][0])
                rows.append([t["text"] for t in current_row])
                # 开始新行
                current_row = [text]
                row_y = text_y
        
        # 处理最后一行
        if current_row:
            current_row.sort(key=lambda x: x["bbox"][0][0])
            rows.append([t["text"] for t in current_row])
            
        return rows

if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description='PDF OCR and Layout Analysis')
    parser.add_argument('pdf_path', help='Path to the PDF file')
    parser.add_argument('--output', '-o', help='Output directory', default=None)
    parser.add_argument('--lang', '-l', help='Language for OCR', default='ch')
    
    args = parser.parse_args()
    
    processor = PDFOCRProcessor(lang=args.lang)
    result_path = processor.process_pdf(args.pdf_path, args.output)
    
    print(f"Processing complete. Results saved to: {result_path}") 